module mod_ode_rk
   use petsc_wrapper
   use mod_ode_solver_abstract
   use stdlib_optval
   use iso_fortran_env, only: dp => real64

   private
   public :: type_ode_rk

   type, extends(type_ode_solver_abstract) :: type_ode_rk

      real(dp), allocatable :: alpha(:), beta(:)

   contains

      procedure :: initialize
      procedure :: destroy
      #: for i in range(1, 5)
      procedure :: step_array${i}$D
      procedure :: substep_array${i}$D
      #: endfor
      procedure :: step_vector
      procedure :: substep_vector

      procedure, private :: time_advance => TimeAdvanceRK
      #: for i in range(1, 5)
      procedure, private :: time_advance_${i}$d => TimeAdvanceRK${i}$D
      #: endfor
      procedure, private :: time_advance_vector => TimeAdvanceRKVector

   end type type_ode_rk

contains

   subroutine initialize(this, scheme)
      class(type_ode_rk), intent(inout) :: this
      character(len=*), intent(in), optional :: scheme
      this%name = "Runge-Kutta"
      this%scheme = optval(scheme, "tvd3")

      select case (this%scheme)
      case ("tvd2")
         call petsc_print("RK scheme: TVD2")
         this%order = 2
         allocate (this%alpha(2))
         allocate (this%beta(2))
         this%alpha(:) = [0.0_dp, 1.0_dp/2.0_dp]
         this%beta(:) = [1.0_dp, 1.0_dp/2.0_dp]
      case ("tvd3")
         call petsc_print("RK scheme: TVD3")
         this%order = 3
         allocate (this%alpha(3))
         allocate (this%beta(3))
         this%alpha(:) = [0.0_dp, 3.0_dp/4.0_dp, 1.0_dp/3.0_dp]
         this%beta(:) = [1.0_dp, 1.0_dp/4.0_dp, 2.0_dp/3.0_dp]
      case default
         call petsc_error_print("unknown rk scheme "//this%scheme//".")
      end select
   end subroutine initialize

   subroutine destroy(this)
      class(type_ode_rk), intent(inout) :: this

      deallocate (this%alpha, this%beta)

   end subroutine destroy

   #: for i, dimension in list(zip(range(1, 5), [":", ":, :", ":, :, :", ":, :, :, :"]))
   subroutine step_array${i}$D(this, solution, ctx)
      class(type_ode_rk), intent(inout) :: this
      real(dp), intent(inout), target :: solution(${dimension}$)
      type(c_ptr), intent(inout) :: ctx
      integer :: lrk
      real(dp) :: t

      this%solution_array${i}$_n(${dimension}$) = solution(${dimension}$)
      this%solution_array${i}$ => solution
      do lrk = 1, this%order
         call this%solve_rhs(t, solution, ctx)
         call this%time_advance(lrk)
      end do

   end subroutine step_array${i}$D

   subroutine substep_array${i}$D(this, lrk, solution, ctx)
      class(type_ode_rk), intent(inout) :: this
      real(dp), intent(inout), target :: solution(${dimension}$)
      type(c_ptr), intent(inout) :: ctx
      integer, intent(in) :: lrk
      real(dp) :: t

      if(lrk == 1) then
         this%solution_array${i}$_n(${dimension}$) = solution(${dimension}$)
      endif
      this%solution_array${i}$ => solution
      call this%solve_rhs(t, solution, ctx)
      call this%time_advance(lrk)

   end subroutine substep_array${i}$D

   #: endfor

   subroutine step_vector(this, solution, ctx)
      class(type_ode_rk), intent(inout) :: this
      class(type_domain_vector), intent(inout), target :: solution
      type(c_ptr), intent(inout) :: ctx
      integer :: lrk
      real(dp) :: t

      if (.not. this%solution_vector_n%allocated) call this%solution_vector_n%duplicate(solution)
      call this%solution_vector_n%copy(solution)
      this%solution_vector => solution
      do lrk = 1, this%order
         call this%solve_rhs(t, solution, ctx)
         call this%time_advance(lrk)
      end do

   end subroutine step_vector

   subroutine substep_vector(this, lrk, solution, ctx)
      class(type_ode_rk), intent(inout) :: this
      class(type_domain_vector), intent(inout), target :: solution
      type(c_ptr), intent(inout) :: ctx
      integer, intent(in) :: lrk
      real(dp) :: t

      if(lrk == 1) then
         if (.not. this%solution_vector_n%allocated) call this%solution_vector_n%duplicate(solution)
         call this%solution_vector_n%copy(solution)
      endif

      this%solution_vector => solution
      call this%solve_rhs(t, solution, ctx)
      call this%time_advance(lrk)

   end subroutine substep_vector

   subroutine TimeAdvanceRK(this, lrk)
      class(type_ode_rk), intent(inout) :: this
      integer, intent(in) :: lrk

      select case (this%data_type)
         #: for i in range(1, 5)
      case ("array${i}$D")
         call this%time_advance_${i}$d(lrk)
         #: endfor
      case ("vector")
         call this%time_advance_vector(lrk)
      end select

   end subroutine TimeAdvanceRK

   #: for i in range(1, 5)
   subroutine TimeAdvanceRK${i}$D(this, lrk)
      class(type_ode_rk), intent(inout) :: this
      integer, intent(in) :: lrk

      this%solution_array${i}$ = this%alpha(lrk)*this%solution_array${i}$_n &
                                 + this%beta(lrk)*(this%solution_array${i}$+this%rhs_${i}$*this%dt)

   end subroutine TimeAdvanceRK${i}$D
   #: endfor
   subroutine TimeAdvanceRKVector(this, lrk)
      class(type_ode_rk), intent(inout) :: this
      integer, intent(in) :: lrk

      call this%solution_vector%axpbypcz(this%alpha(lrk), this%beta(lrk)*this%dt, this%beta(lrk), &
                                         this%solution_vector_n, this%rhs_vector)
   end subroutine TimeAdvanceRKVector

end module mod_ode_rk
